/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "ola_speke.h"
#include "ola_error.h"
#include "ola_log.h"
#include "ola_pal_os.h"
#include "ola_negotiation.h"
#include "ola_base_utils.h"
#include "session_key_agree.h"
#include "securec.h"
#include "ola_security.h"
#include "ola_sdk_api.h"
#include "ola_random.h"

#define SPEKE_NEGO_TIMEOUT (15 * 1000)

#define OLA_SPEKE_HANDLE "OLA"

#define SPEKE_SESSION_ID_LEN    32
#define SESSIONKEY_LEN          64
#define SPEKE_JSON_MAX_LEN      2048
#define AES256_KEY_BIT_LENGTH   256

#define SPEKE_STRING_SESSION_ID     "sessionId"
#define SPEKE_STRING_SECURITY_DATA  "securityData"

typedef enum {
    OLA_SPEKE_STEP_FIRST,
    OLA_SPEKE_STEP_FINISH
} OlaSpekeStep;

typedef struct {
    char sessionId[SPEKE_SESSION_ID_LEN + 1];
    unsigned char sessionKey[SESSIONKEY_LEN + 1];
    unsigned long long lastMsgTime;
} OlaSpekeInfo;

static SKA_Callback g_sessionCb = {0};
static OlaNegotiationCb g_negoCb = {NULL, NULL};
static OlaSpekeInfo g_spekeInfo = {0};

static OlaNegoStatus g_negoStatus = OLA_NEGO_STATUS_UNINIT;

static void SetSpekeNegoStatus(OlaNegoStatus status)
{
    OlaLogInfo("set speke nego status[%d]", status);
    g_negoStatus = status;
}

static OlaNegoStatus SpekeGetNegoStatus(void)
{
    return g_negoStatus;
}

/* 重置speke协商状态 */
static void ResetSpekeStatus(void)
{
    (void)memset_s(&g_spekeInfo, sizeof(OlaSpekeInfo), 0, sizeof(OlaSpekeInfo));
    SetSpekeNegoStatus(OLA_NEGO_STATUS_INIT);
}

/* 字符数组转16进制字符串 */
static void ByteToHexStr(const uint8_t* source, char* dest, uint32_t sourceLen)
{
    if ((source == NULL) || (dest == NULL)) {
        OlaLogError("invalid param");
        return;
    }
    uint32_t i;
    unsigned char highByte, lowByte;

    for (i = 0; i < sourceLen; i++) {
        highByte = source[i] >> 4;
        lowByte = source[i] & 0x0f;

        highByte += 0x30;

        if (highByte > 0x39) {
            dest[i * 2] = highByte + 0x07;
        } else {
            dest[i * 2] = highByte;
        }
        lowByte += 0x30;
        if (lowByte > 0x39) {
            dest[i * 2 + 1] = lowByte + 0x07;
        } else {
            dest[i * 2 + 1] = lowByte;
        }
    }
    return;
}

/* 协商session过程中获取PIN */
static int GetPinCode(char *buffer, int bufferLen)
{
    if ((buffer == NULL) || (bufferLen == 0)) {
        OlaLogError("invalid param");
        return OLA_ERROR;
    }
    if (g_negoCb.getPinCodeCb == NULL) {
        OlaLogError("pin code callback null");
        return OLA_ERROR;
    }
    int len = bufferLen;
    if (g_negoCb.getPinCodeCb((unsigned char *)buffer, (unsigned int *)&len) != OLA_OK) {
        OlaLogError("get pinCode failed");
        return OLA_ERROR;
    }
    if (len > bufferLen) {
        OlaLogError("pinCode failed");
        return OLA_ERROR;
    }
    OlaLogInfo("get pincode SUCCESS");
    return len;
}

/* session协商完成回调 */
static void OnSessionAgreed(void *clientHandle, const uint8_t *sessionKey, uint32_t sessionKeyLen)
{
    OlaLogDebug("session agreed in");
    if ((sessionKey == NULL) || (sessionKeyLen == 0)) {
        OlaLogError("invalid param, keyLen[%u]", sessionKeyLen);
        return;
    }
    (void)clientHandle;
    char dstKey[SESSIONKEY_LEN + 1] = {0};
    ByteToHexStr(sessionKey, dstKey, sessionKeyLen);
    (void)memset_s(g_spekeInfo.sessionKey, SESSIONKEY_LEN + 1, 0, SESSIONKEY_LEN + 1);
    if (memcpy_s(g_spekeInfo.sessionKey, sizeof(g_spekeInfo.sessionKey), sessionKey, sessionKeyLen) != OLA_OK) {
        OlaLogError("get sessionKey fail\n");
        return;
    }
    OlaLogInfo("get sessionKey success");
    SetSpekeNegoStatus(OLA_NEGO_STATUS_FINISH);
}

static void OnKeyAgreeFinish(void *clientHandle)
{
    (void)clientHandle;
    OlaLogInfo("Speke Finish");
}

static void OnKeyAgreeError(void *clientHandle, int errorCode)
{
    OlaLogDebug("session agreed error in");
    (void)clientHandle;
    OlaLogError("session failed[%d]", errorCode);
    /* 重置状态，等待下一次协商 */
    ResetSpekeStatus();
}

/* 协商过程中，给hichain调用的接口，将处理好的数据发出去 */
static int TransmitSessionData(void *clientHandle, const uint8_t *data, uint32_t dataLen)
{
    OlaLogDebug("transmit session data in");
    (void)clientHandle;
    if ((data == NULL) || (dataLen == 0) || (data[dataLen - 1] != 0)) {
        OlaLogError("invalid param");
        return ERR_SKA_INVALID_PARAM;
    }

    cJSON *json = cJSON_CreateObject();
    if (json == NULL) {
        OlaLogError("create json error");
        return ERR_SKA_FAILED;
    }
    int ret = ERR_SKA_FAILED;
    do {
        if (cJSON_AddStringToObject(json, SPEKE_STRING_SESSION_ID, (const char *)(g_spekeInfo.sessionId)) == NULL) {
            OlaLogError("add string error");
            break;
        }
        cJSON *dataJson = cJSON_Parse((const char *)data);
        if (dataJson == NULL) {
            OlaLogError("parse json error");
            break;
        }

        if (cJSON_AddItemToObject(json, SPEKE_STRING_SECURITY_DATA, dataJson) == false) {
            OlaLogError("add object error");
            break;
        }

        char *payload = cJSON_PrintUnformatted(json);
        if (payload == NULL) {
            OlaLogError("json print error");
            break;
        }
        OlaData sendData = {(unsigned char *)payload, strlen(payload)};
        if (g_negoCb.sendData(&sendData) != OLA_OK) {
            OlaLogError("send data error");
            OlaPutNullAfterFree((void**)&payload);
            break;
        }
        ret = (int)strlen(payload);
        OlaLogInfo("send response success");
        OlaPutNullAfterFree((void**)&payload);
    } while (0);

    cJSON_Delete(json);
    return ret;
}

static void SessionCallbackInit(void)
{
    g_sessionCb.getPinCode = GetPinCode;
    g_sessionCb.onSessionKeyAgreed = OnSessionAgreed;
    g_sessionCb.onKeyAgreeFinish = OnKeyAgreeFinish;
    g_sessionCb.onKeyAgreeError = OnKeyAgreeError;
    g_sessionCb.transmitSessionData = TransmitSessionData;
}

static SKA_Callback *GetSessionCallback(void)
{
    return &g_sessionCb;
}

static int SpekeInit(const OlaNegotiationCb *negoCb)
{
    /* 处理入参回调函数 */
    if ((negoCb == NULL) || (negoCb->sendData == NULL) || (negoCb->getPinCodeCb == NULL)) {
        OlaLogError("invalid cb param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    g_negoCb.sendData = negoCb->sendData;
    g_negoCb.getPinCodeCb = negoCb->getPinCodeCb;
    /* 初始化Speke相关，如session钩子 */
    SessionCallbackInit();
    int ret = SKA_InitKeyAgreeSdk(false, OLA_SPEKE_HANDLE);
    if (ret != ERR_SKA_SUCCESS) {
        OlaLogError("init speke failed2, ret[%d]", ret);
        return OLA_ERROR;
    }
    SetSpekeNegoStatus(OLA_NEGO_STATUS_INIT);

    return OLA_OK;
}

static int UpdateSpekeInfo(const OlaData *sessionId)
{
    if ((sessionId != NULL) && (sessionId->data != NULL) && (sessionId->len != 0)) {
        if (strncpy_s(g_spekeInfo.sessionId, sizeof(g_spekeInfo.sessionId),
            (const char *)sessionId->data, sessionId->len) != OLA_OK) {
            OlaLogError("copy sessionId failed");
            return OLA_SECUREC_ERR_STRNCPY;
        }
    }

    int ret = OlaGetTime(&g_spekeInfo.lastMsgTime);
    if (ret != OLA_OK) {
        OlaLogError("update time faied, ret[%d]", ret);
        return ret;
    }

    return OLA_OK;
}

static int GetSpekeMsgStep(const OlaData *sessionId, int *msgStep)
{
    if ((sessionId == NULL) || (sessionId->data == NULL) || (msgStep == NULL)) {
        OlaLogError("session id invalid");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    bool isTimeout = true;
    if (OlaIsOvertime(&isTimeout, g_spekeInfo.lastMsgTime, SPEKE_NEGO_TIMEOUT) != OLA_OK) {
        OlaLogError("speke,get timeout state fail");
        return OLA_ERROR;
    }

    if ((strlen(g_spekeInfo.sessionId) == 0) || isTimeout) {
        *msgStep = OLA_SPEKE_STEP_FIRST;
    } else if (strcmp(g_spekeInfo.sessionId, (char *)sessionId->data) == OLA_OK) {
        *msgStep = OLA_SPEKE_STEP_FINISH;
    }

    OlaLogInfo("timeout flag: %d, step: %d", isTimeout, *msgStep);
    return OLA_OK;
}

static int GetSessionIdAndSecurityData(const OlaData *inData, OlaData *sessionId)
{
    if ((inData == NULL) || (inData->len == 0) || (inData->data == NULL)) {
        OlaLogError("invalid speke data");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    if ((sessionId == NULL)) {
        OlaLogError("out buffer null");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    cJSON *spekeJson = cJSON_Parse((char *)inData->data);
    if (spekeJson == NULL) {
        OlaLogError("parse msg failed");
        return OLA_CJSON_ERR_PARSE;
    }
    sessionId->data = (unsigned char *)OlaGetJsonString(spekeJson, OLA_STRING_SESSIONID);
    cJSON_Delete(spekeJson);
    if (sessionId->data == NULL) {
        OlaLogError("get sessionId failed");
        return OLA_ERROR;
    }
    sessionId->len = strlen((char *)sessionId->data);
    return OLA_OK;
}

static int HandleFirstMsg(const char *spekeBuf, unsigned int bufLen, const OlaData *sessionId)
{
    if ((spekeBuf == NULL) || (sessionId == NULL) || (sessionId->data == NULL)) {
        OlaLogError("param null");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    if ((bufLen == 0) || (sessionId->len == 0)) {
        OlaLogError("invalid len, bufLen[%u], sessionId len[%u]", bufLen, sessionId->len);
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    ResetSpekeStatus();

    if (UpdateSpekeInfo(sessionId) != OLA_OK) {
        OlaLogError("update speke info failed");
        return OLA_ERROR;
    }

    /* 处理初次包数据 */
    SKA_Chars firstData = {NULL, 0};
    firstData.chars = OlaMallocAndCopyStr(spekeBuf, bufLen);
    firstData.length = strlen(spekeBuf);
    if (firstData.chars == NULL) {
        OlaLogError("copy speke msg failed");
        return OLA_ERROR;
    }
    int ret = SKA_StartKeyAgreeSession(OLA_SPEKE_HANDLE, &firstData, GetSessionCallback());
    if (ret != ERR_SKA_SUCCESS) {
        OlaLogError("handle speke msg failed[%d]", ret);
        if (firstData.chars != NULL) {
            OlaFree(firstData.chars);
        }
        ResetSpekeStatus();
        return OLA_COAP_RESP_NOT_RESP;
    }

    SetSpekeNegoStatus(OLA_NEGO_STATUS_DOING);
    if (firstData.chars != NULL) {
        OlaFree(firstData.chars);
    }
    return OLA_COAP_RESP_NOT_RESP;
}

static int HandleSecondMsg(const char *spekeBuf, unsigned int bufLen)
{
    if (SpekeGetNegoStatus() != OLA_NEGO_STATUS_DOING) {
        OlaLogError("now nego status[%d], handle secone msg failed", SpekeGetNegoStatus());
        return OLA_ERROR;
    }
    if (UpdateSpekeInfo(NULL) != OLA_OK) {
        OlaLogError("update speke info failed");
        return OLA_ERROR;
    }
    /* 处理第二包数据 */
    SKA_Chars secondData = {NULL, 0};
    secondData.chars = OlaMallocAndCopyStr(spekeBuf, bufLen);
    secondData.length = strlen(spekeBuf);
    if (secondData.chars == NULL) {
        OlaLogError("copy speke msg failed");
        return OLA_ERROR;
    }
    int ret = SKA_ProcessKeyAgreeSessionData(OLA_SPEKE_HANDLE, &secondData);
    if (ret != ERR_SKA_SUCCESS) {
        OlaLogError("handle speke msg failed [%d]", ret);
        if (secondData.chars != NULL) {
            OlaFree(secondData.chars);
        }
        return OLA_COAP_RESP_NOT_RESP;
    }
    SetSpekeNegoStatus(OLA_NEGO_STATUS_FINISH);
    if (secondData.chars != NULL) {
        OlaFree(secondData.chars);
    }
    return OLA_COAP_RESP_NOT_RESP;
}

static int SpekeProcessData(const OlaData *inData, OlaData *outData)
{
    /* speke协商不在此接口回复信息，因此outData不使用 */
    (void)outData;
    if ((inData == NULL) || (inData->data == NULL) || (inData->len == 0)) {
        OlaLogError("invalid param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }
    OlaData sessionId = {NULL, 0};
    if ((GetSessionIdAndSecurityData(inData, &sessionId) != OLA_OK) || (sessionId.len != SPEKE_SESSION_ID_LEN)) {
        OlaLogError("parse speke buf faild, session id len[%u]", sessionId.len);
        return OLA_ERROR;
    }

    int ret;
    /* 判断接收到的协商信息是第几包 */
    int msgStep = -1;
    if (GetSpekeMsgStep(&sessionId, &msgStep) != OLA_OK) {
        OlaLogError("get msg step failed");
        return OLA_ERROR;
    }
    if (msgStep == OLA_SPEKE_STEP_FIRST) {
        ret = HandleFirstMsg((const char *)inData->data, inData->len, &sessionId);
    } else if (msgStep == OLA_SPEKE_STEP_FINISH) {
        ret = HandleSecondMsg((const char *)inData->data, inData->len);
    } else {
        OlaLogError("get msg step failed");
        ret = OLA_ERROR;
    }

    OlaLogInfo("speke data process, msgStep[%d], ret[%d]", msgStep, ret);
    return ret;
}

/* 检查加解密入参数据是否正确 */
static int CheckDataParamAndStatus(const OlaData *inData, const OlaData *outData, bool isDecode)
{
    if ((inData == NULL) || (inData->data == NULL) || (inData->len == 0)) {
        OlaLogError("in data invalid");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    if (isDecode && (inData->len < (OLA_TYPICAL_GCM_IV_LEN + OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_TAG_LEN))) {
        OlaLogError("in data len invalid, inLen[%u]", inData->len);
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    if ((outData == NULL) || (outData->data != NULL) || (outData->len != 0)) {
        OlaLogError("out data invalid");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    if (SpekeGetNegoStatus() != OLA_NEGO_STATUS_FINISH) {
        OlaLogError("negotiation is unfinished");
        return OLA_ERROR;
    }

    return OLA_OK;
}

static int SpekeEncodeData(const OlaData *inData, OlaData *outData)
{
    if (CheckDataParamAndStatus(inData, outData, false) != OLA_OK) {
        return OLA_ERROR;
    }
    OlaData base64 = {NULL, 0};
    /* TODO speke加密报文待定，当前|ver(1byte)|iv(12bytes)|data|tag(16bytes)|且为aes128gcm加密 */
    base64.len = inData->len + OLA_TYPICAL_GCM_VERSION_LEN +
        OLA_TYPICAL_GCM_IV_LEN + OLA_TYPICAL_GCM_TAG_LEN;
    base64.data = OlaMalloc(base64.len);
    if (base64.data == NULL) {
        OlaLogError("malloc error");
        return OLA_OS_ERR_MALLOC;
    }
    base64.data[0] = OLA_TYPICAL_GCM_VERSION;
    int ret;
    do {
        ret = OlaRand(base64.data + OLA_TYPICAL_GCM_VERSION_LEN, OLA_TYPICAL_GCM_IV_LEN);
        if (ret != OLA_OK) {
            OlaLogError("gen iv failed");
            break;
        }

        AesGcmParam inputParam = {
            .key = {g_spekeInfo.sessionKey, OLA_AES_128_GCM_KEY_LEN},
            .iv = {base64.data + OLA_TYPICAL_GCM_VERSION_LEN, OLA_TYPICAL_GCM_IV_LEN},
            .add = {NULL, 0},
            .tag = {NULL, 0},
            .inData = {inData->data, inData->len},
        };

        OlaBuffer encData = {base64.data + OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_IV_LEN, 0, inData->len};
        OlaBuffer tag = {base64.data + OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_IV_LEN + inData->len,
            0, OLA_TYPICAL_GCM_TAG_LEN};
        ret = OlaAesGcmEncript(&inputParam, &encData, &tag);
        if (ret != OLA_OK) {
            OlaLogWarning("aes-128-gcm encript error");
            break;
        }
        ret = OlaEncodeBase64Data(&base64, outData);
        if (ret != OLA_OK) {
            OlaLogWarning("base 64 encode error");
            break;
        }
        OlaDataFree(&base64);
        return OLA_OK;
    } while (0);

    OlaDataFree(outData);
    OlaDataFree(&base64);
    return ret;
}

static int SpekeDecodeData(const OlaData *inData, OlaData *outData)
{
    if (CheckDataParamAndStatus(inData, outData, true) != OLA_OK) {
        return OLA_ERROR;
    }
    OlaData base64Data = {NULL, 0};
    int ret = OlaDecodeBase64Data(inData, &base64Data);
    if (ret != OLA_OK) {
        OlaLogError("base64 decode data failed, ret[%d]", ret);
        return ret;
    }

    /* TODO 解密组包形式未定，当前|ver(1byte)|iv(12bytes)|data|tag(16bytes)|且为aes128gcm解密 */
    AesGcmParam inputParam = {
        .key = {g_spekeInfo.sessionKey, OLA_AES_128_GCM_KEY_LEN},
        .iv = {base64Data.data + OLA_TYPICAL_GCM_VERSION_LEN, OLA_TYPICAL_GCM_IV_LEN},
        .add = {NULL, 0},
        .tag = {base64Data.data + base64Data.len - OLA_TYPICAL_GCM_TAG_LEN, OLA_TYPICAL_GCM_TAG_LEN},
        .inData = {
            base64Data.data + OLA_TYPICAL_GCM_VERSION_LEN + OLA_TYPICAL_GCM_IV_LEN,
            base64Data.len - OLA_TYPICAL_GCM_VERSION_LEN - OLA_TYPICAL_GCM_IV_LEN - OLA_TYPICAL_GCM_TAG_LEN
        },
    };

    outData->data = OlaMalloc(inputParam.inData.len + 1);
    if (outData->data == NULL) {
        OlaPutNullAfterFree((void**)&base64Data.data);
        OlaLogError("malloc error");
        return OLA_OS_ERR_MALLOC;
    }
    outData->data[inputParam.inData.len] = 0;
    OlaBuffer outBuffer = {outData->data, 0, inputParam.inData.len + 1};
    ret = OlaAesGcmDecript(&inputParam, &outBuffer);
    if (ret != OLA_OK) {
        OlaDataFree(outData);
        outBuffer.buf = NULL;
    } else {
        outData->len = outBuffer.len;
    }
    OlaDataFree(&base64Data);
    return ret;
}

static void SpekeDeinit(void)
{
    /* 重置所有内容，如sessionKey、状态、session回调函数 */
    SetSpekeNegoStatus(OLA_NEGO_STATUS_UNINIT);
    (void)memset_s(&g_spekeInfo, sizeof(OlaSpekeInfo), 0, sizeof(OlaSpekeInfo));

    g_sessionCb.getPinCode = NULL;
    g_sessionCb.onSessionKeyAgreed = NULL;
    g_sessionCb.onKeyAgreeFinish = NULL;
    g_sessionCb.onKeyAgreeError = NULL;
    g_sessionCb.transmitSessionData = NULL;

    g_negoCb.sendData = NULL;
    g_negoCb.getPinCodeCb = NULL;
    SKA_ReleaseKeyAgreeSdk();
}

static OlaNegotiationHook g_spekeNegoHook = {
    .init = SpekeInit,
    .processData = SpekeProcessData,
    .getStatus = SpekeGetNegoStatus,
    .encodeData = SpekeEncodeData,
    .decodeData = SpekeDecodeData,
    .deinit = SpekeDeinit,
};

const OlaNegotiationHook *OlaGetSpekeNegoHooks(void)
{
    return &g_spekeNegoHook;
}
